//
// Created by lxinp on 2022/10/18.
//

#include "avl_tree.h"


template<typename K, typename V>
AVLTree<K, V>::AVLTree() : root(nullptr) {}


template<typename K, typename V>
AVLTree<K, V>::AVLTree(const AVLTree& other) : root(copy(other.root)) {}


template<typename K, typename V>
AVLTree<K, V>::~AVLTree() {
    clear(root);
    root = nullptr;
}


template<typename K, typename V>
AVLTree<K, V>& AVLTree<K, V>::operator=(const AVLTree& other) {
    if (this != &other) {
        if (root) clear(root);
        root = copy(other.root);
    }
    return *this;
}


template<typename K, typename V>
void AVLTree<K, V>::clear() {
    clear(root);
}


template<typename K, typename V>
void AVLTree<K, V>::insert(const K& key, const V& value) {
    insert(root, key, value);
}


template<typename K, typename V>
void AVLTree<K, V>::remove(const K& key) {
    remove(root, key);
}


template<typename K, typename V>
V AVLTree<K, V>::find(const K &key) const {
    return find(root, key);
}


template<typename K, typename V>
typename AVLTree<K, V>::Node* AVLTree<K, V>::copy(const AVLTree::Node* other) {
    if (!other) return nullptr;
    Node* node = new Node(other->key, other->value);
    node->left = copy(other->left);
    node->right = copy(other->right);
    return node;
}


template<typename K, typename V>
void AVLTree<K, V>::clear(AVLTree::Node* subtree) {
    if (!subtree) return;
    clear(subtree->left);
    clear(subtree->right);
    delete subtree;
    subtree = nullptr;
}


template<typename K, typename V>
void AVLTree<K, V>::insert(AVLTree::Node* subtree, const K& key, const V& value) {
    if (!subtree) subtree = new Node(key, value);
    else {
        if (key < subtree->key) insert(subtree->left, key, value);
        else insert(subtree->right, key, value);
    }
    rebalance(subtree);
}


template<typename K, typename V>
void AVLTree<K, V>::remove(AVLTree::Node*& subtree, const K& key) {
    if (!subtree) return;

    if (key < subtree->key) {
        remove(subtree->left, key);
        rebalance(subtree);
    } else if (key > subtree->key) {
        remove(subtree->right, key);
        rebalance(subtree);
    } else {
        // Reached the node to remove
        if (!subtree->left && !subtree->right) {
            delete subtree;
            subtree = nullptr;
            return;
        } else if (subtree->left && subtree->right) {
            // Go to the rightmost node of the left child
            // Swap with the current node
            // Delete that node
            Node* succ = subtree->left;
            Node* pre = succ;
            while (succ && succ->right) succ = succ->right;
            while (pre->right != succ) pre = pre ->right;
            pre->right = succ->left;
            swap(subtree, succ);
            delete succ;
        } else {
            // Set the current node to the only child
            Node* temp = subtree;
            subtree = std::max(subtree->left, subtree->right);
            delete temp;
        }
        rebalance(subtree);
    }
}


template<typename K, typename V>
V AVLTree<K, V>::find(AVLTree::Node* subtree, const K& key) const {
    if (!subtree) return V();
    else if (subtree->key == key) return subtree->value;
    else if (subtree->key < key) return find(subtree->right, key);
    else return find(subtree->left, key);
}


/**
 *  A <- node
 * / \
 *    B <- temp
 *   / \
 *      O
 */
template<typename K, typename V>
void AVLTree<K, V>::rotate_left(AVLTree::Node*& node) {
    Node* temp = node->right;
    node->right = temp->left;
    temp->left = node;
    update_height(temp);
    update_height(node);
    node = temp;
}


/**
 *     A <- node
 *    / \
 *   B <- temp
 *  / \
 * O
 */
template<typename K, typename V>
void AVLTree<K, V>::rotate_right(AVLTree::Node*& node) {
    Node* temp = node->left;
    node->left = temp->right;
    temp->right = node;
    update_height(temp);
    update_height(node);
    node = temp;
}


/**
 *    A <- Node
 *   /
 *  B
 *   \
 *    O
 */
template<typename K, typename V>
void AVLTree<K, V>::rotate_left_right(AVLTree::Node*& node) {
    rotate_left(node->left);
    rotate_right(node);
}


/**
 * A <- node
 *  \
 *   B
 *  /
 * O
 */
template<typename K, typename V>
void AVLTree<K, V>::rotate_right_left(AVLTree::Node*& node) {
    rotate_right(node->right);
    rotate_left(node);
}


/**
 * 5 cases:
 *  A <- node   |     A <- node
 * / \          |    / \
 *    B <- temp |   B <- temp
 *   / \        |  / \
 *      O       | O
 * ============================
 *   A <- Node  |  A <- node
 *  /           |   \
 * B            |    B
 *  \           |   /
 *   O          |  O
 * And a balanced case (still need to update height!)
 */
template<typename K, typename V>
void AVLTree<K, V>::rebalance(AVLTree::Node*& subtree) {
    if (height(subtree->left) - height(subtree->right) == 2) {
        Node* left = subtree->left;
        if (height(left->left) - height(left->right) == 1) rotate_right(subtree);
        else if (height(left->left) - height(left->right) == -1) rotate_left_right(subtree);
    }
    if (height(subtree->right) - height(subtree->left) == 2) {
        Node* right = subtree->right;
        if (height(right->right) - height(right->left) == 1) rotate_left(subtree);
        else if (height(right->right) - height(right->left) == -1) rotate_right_left(subtree);
    }
    update_height(subtree);
}


template<typename K, typename V>
int AVLTree<K, V>::height(const AVLTree::Node* subtree) const {
    if (!subtree) return -1;  // return 0?
    return subtree->height;
}


template<typename K, typename V>
void AVLTree<K, V>::update_height(AVLTree::Node* subtree) {
    subtree->height = 1 + std::max(height(subtree->left), height(subtree->right));
}


template<typename K, typename V>
void AVLTree<K, V>::swap(AVLTree::Node*& first, AVLTree::Node*& second) {
    K temp_key = first->key;
    V temp_val = first->value;
    first->key = second->key;
    first->value = second->value;
    second->key = temp_key;
    second->value = temp_val;
}
